////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//                                                                                                                            //
//                                                         WNProject                                                          //
//                                                                                                                            //
//            This file is distributed under the BSD 2-Clause open source license. See License.txt for details.               //
//                                                                                                                            //
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#include "WNCore/inc/WNPlatform.h"
#include "WNCore/inc/WNTypes.h"
#include "WNCore/inc/WNAtomic.h"
#include "WNConcurrency/inc/WNRecursiveMutex.h"
#include "WNFibers/src/WNFibersPrivate.h"
#include "WNFibers/src/WNFiberPrivateData.h"
#include "WNFibers/src/WNFiberManagerImpl.h"
#include "WNMemory/inc/WNMemory.h"

#ifdef _WN_X86
    extern "C" {
        WN_VOID WN_CDECL_BEGIN BeginFiber(WN_POINTER retptrloc, WN_POINTER oldstackloc, WN_POINTER inValue, WN_POINTER func, WN_POINTER stackloc, WN_UINT32 stacksize, WN_UINT32 stackleft) WN_CDECL_END;
        WN_VOID WN_CDECL_BEGIN EndFiber(WN_POINTER pushstacklocation, WN_POINTER oldscriptpointer, WN_SIZE_T cont, WN_POINTER retpt, WN_SIZE_T* setrunnable, WN_POINTER* retloc) WN_CDECL_END;
        WN_VOID WN_CDECL_BEGIN SwitchFiber(WN_POINTER* pushstacklocation, WN_POINTER oldscriptptr, WN_SIZE_T cont, WN_POINTER retpt, WN_SIZE_T* setrunnable, WN_POINTER* retloc, WN_POINTER inValue, WN_UINT32 stacksize, WN_UINT32 stackleft) WN_CDECL_END;
    }
#endif

using namespace WNCore;
using namespace WNConcurrency;
using namespace WNFibers;
using namespace WNMemory;

static WN_THREAD_LOCAL WN_BOOL tlContinuousRunning = WN_FALSE;
static WN_THREAD_LOCAL WNFiberData* tlCurStack = 0;

#ifdef _WN_FIBER_USE_LIBRARIES
    static WN_THREAD_LOCAL WNFiberData* tlMainStack = 0;
    static WN_THREAD_LOCAL WNFiberData* tlLastStack = 0;
    static WN_THREAD_LOCAL WN_UINT32 tlLastStackCont = 0;
#else
    static WN_THREAD_LOCAL WN_POINTER tlRetStackPtr = 0;
    static WN_THREAD_LOCAL WN_POINTER tlRetPtr = 0;
#endif

WN_FORCE_INLINE WNFiberData* GetRunningFiber() {
    return(tlCurStack);
}

//return to the owning location (jump off the secondary stack back onto the main stack)
WN_FORCE_INLINE WN_VOID ReturnToMainStack(WN_UINT32 _continue, WNFiberManager* _manager) {
    #ifdef _WN_FIBER_USE_LIBRARIES
        _continue = (_continue == 0) ? 0 : 1;

        WNFiberData* oldData = tlCurStack;

        if (tlContinuousRunning) {
            WNFiberData* next = _manager->GetNextStack();

            if (next != WN_NULL) {
                tlCurStack = next;
                tlLastStack = oldData;
                tlLastStackCont = _continue;

                WN_SWITCH_FIBERS(oldData->privateData->fiberStruct, tlCurStack->privateData->fiberStruct);

                return;
            }
        }

        tlCurStack = tlMainStack;
        tlLastStack = oldData;
        tlLastStackCont = _continue;

        WN_SWITCH_FIBERS(oldData->privateData->fiberStruct, tlCurStack->privateData->fiberStruct);
    #else
        WNFiberData* si = tlCurStack;

        _continue = (_continue == 0) ? 0 : 1; //clamp ret to 0 or 1

        //return from whence we came
        if (tlContinuousRunning) {
            WNFiberData* next = _manager->GetNextStack();

            if (next != WN_NULL) {
                const WN_INT32 stackleft = static_cast<WN_INT32>(next->privateData->newstackloc - next->privateData->stack);

                if (stackleft < 0) {
                    WN_PRINTF("ERROR YOU BLEW YOUR STACK\n");

                    return;
                }

                tlCurStack = next;

                SwitchFiber(&si->privateData->newstackloc, next->privateData->newstackloc, _continue, next->privateData->returnstacklocation, &si->runnable, &si->privateData->returnstacklocation, si->startVal, next->privateData->stacksize, stackleft);

                return;
            }
        }

        EndFiber((WN_POINTER)&si->privateData->newstackloc, (WN_POINTER)tlRetStackPtr, _continue, (WN_POINTER)tlRetPtr, &si->runnable, &si->privateData->returnstacklocation);
    #endif
}

namespace WNFibers {
    WN_VOID WNFiberYield() {
        WNFiberData* si = GetRunningFiber();

        si->sleeptype = WN_SLP_NO_SLEEP;

        ReturnToMainStack(1, si->privateData->manager);

        #ifdef _WN_FIBER_USE_LIBRARIES
            tlLastStack->privateData->cont = tlLastStackCont;
        #endif
    }

    WN_VOID WNFiberSleepWithAlarm(WN_UINT32 _alarmticks) {
        WNFiberData* si = GetRunningFiber();

        si->sleepTicks = _alarmticks;
        si->sleeptype = WN_SLP_COUNTED_SIGNAL;

        ReturnToMainStack(1, si->privateData->manager);

        #ifdef _WN_FIBER_USE_LIBRARIES
            tlLastStack->privateData->cont = tlLastStackCont;
        #endif
    }

    WN_VOID WNFiberSleep(WN_UINT32 _time) {
        WNFiberData* si = GetRunningFiber();

        WN_UINT32 t = si->privateData->manager->GetTime();

        t += _time;

        si->sleepval = t;
        si->sleeptype = WN_SLP_SLEEP_TIME;

        ReturnToMainStack(1, si->privateData->manager);

        #ifdef _WN_FIBER_USE_LIBRARIES
            tlLastStack->privateData->cont = tlLastStackCont;
        #endif
    }

    WN_VOID WNFiberReturn(WN_POINTER _value) {
        WNFiberData* si = GetRunningFiber();

        si->sleeptype = WN_SLP_FINISHED;
        si->retVal = _value;

        ReturnToMainStack(0, si->privateData->manager);

        #ifdef _WN_FIBER_USE_LIBRARIES
            tlLastStack->privateData->cont = tlLastStackCont;
        #endif
    }

    WNFiberManager* WNGetFiberManager() {
        return(WN_PLACEMENT_NEW(WNFiberManagerImpl, WNMallocT<WNFiberManagerImpl>()));
    }

    WNFiberManager* WNGetFiberManager(WN_MALLOC_FUNCTION_T _allocator) {
        WN_RELEASE_ASSERT(_allocator != WN_NULL);

        return(WN_PLACEMENT_NEW(WNFiberManagerImpl, _allocator(sizeof(WNFiberManagerImpl))));
    }
}

WNFiberManagerImpl::WNFiberManagerImpl() :
    mThreadUpdating(0),
    mDoneSpin(0),
    mMaxStacks(0),
    mNumActive(0) {
}

WNFiberManagerImpl::~WNFiberManagerImpl() { 
    for (std::vector<WNFiberData*>::iterator i = mThreads.begin(); i != mThreads.end(); ++i){
        WNFree(*i);
    }
}

WN_BOOL WNFiberManagerImpl::Initialize(WN_UINT32 maxStack, WN_BOOL updateQueues, WN_TICK_COUNT_FUNCTION_T timeFunction) {
    mMaxStacks = maxStack;
    mTimeFunction = timeFunction;
    mUpdateQueuesOnDemand = updateQueues;

    return(WN_TRUE);
}

void WNFiberManagerImpl::DestroyFiberData(WNFiberData* _fiberData) {
    WNFree(_fiberData->privateData->stack);
    WNFree(_fiberData->privateData);
    WNFree(_fiberData);
}

WNFiberData* WNFiberManagerImpl::CreateFiberData(WN_FIBER_FUNCTION_T _function, WN_POINTER startval, WN_UINT32 stacksize) {
    WNFiberData * s = WNMallocT<WNFiberData>();

    s->privateData = WNMallocT<WNFiberPrivateData>();
    s->privateData->stack = static_cast<WN_POINTER>(WNMalloc(stacksize));
    s->startVal = startval;
    s->sleeptype = WN_SLP_NOT_STARTED;
    s->privateData->stacksize = stacksize;
    s->runnable = 1;
    s->privateData->manager = this;

    #ifdef _WN_FIBER_USE_LIBRARIES
        WN_ALLOCATE_FIBER(s->privateData->fiberStruct, s->privateData->stack, stacksize);
        WN_INITIALIZE_FIBER(s->privateData->fiberStruct, reinterpret_cast<WN_VOID (*)()>(_function), s->startVal);
    #else
        s->privateData->returnstacklocation = reinterpret_cast<WN_POINTER>(_function);
        s->privateData->newstackloc = s->privateData->stack + s->privateData->stacksize - 16;
        s->privateData->newstackloc = s->privateData->newstackloc - (reinterpret_cast<WN_UINT32>(s->privateData->newstackloc) & 0xFF);
    #endif

    #ifdef _WN_FIBER_VERIFY_STACK
        WNMemSet(s->privateData->stack, 0xAB, stacksize);
    #endif

    return(s);
}

WNFiberData* WNFiberManagerImpl::CreateFiber(WN_FIBER_FUNCTION_T _function, WN_POINTER _startValue, WN_UINT32 _stackSize, WN_POINTER _userData) {
    while (WNAtomicCompareSwap(&mThreadUpdating, 1, 0) != 0) {
        continue;
    }

    mThreadsLock.Lock();

    WNFiberData* fiberData = CreateFiberData(_function, _startValue, _stackSize);

    fiberData->privateData->userdata = _userData;

    mThreads.push_back(fiberData);
    mThreadUpdating = 0;
    mNumActive++;

    mThreadsLock.Unlock();

    mThreadUpdating = 0;

    return(fiberData);
}

WNFiberData *WNFiberManagerImpl::GetNextStack() {
    WNFiberData* s = WN_NULL;

    mRunnableLock.Lock();

    if (mRunnableThreads.size() > 0) {
        s = mRunnableThreads.front();

        mRunnableThreads.pop_front();
    }

    mRunnableLock.Unlock();

    if (s != WN_NULL) {
        s->runnable = 0;
    } else if (mUpdateQueuesOnDemand) {
        UpdateQueues(GetTime());
    }

    return(s);
}

WN_VOID WNFiberManagerImpl::ContinueOther(WNFiberData* s) {
    s->sleeptype = WN_SLP_RUNNING;

    #ifdef _WN_FIBER_USE_LIBRARIES
        WNFiberData* oldData = tlCurStack;

        tlCurStack = s;

        WN_SWITCH_FIBERS(oldData->privateData->fiberStruct, tlCurStack->privateData->fiberStruct);
    #else
        const WN_FIBER_FUNCTION_T func = reinterpret_cast<WN_FIBER_FUNCTION_T>(s->privateData->returnstacklocation);
        const WN_POINTER stackloc = static_cast<WN_POINTER>(s->privateData->newstackloc);
        const WN_INT32 stackleft = static_cast<WN_INT32>(s->privateData->newstackloc - s->privateData->stack );

        if (stackleft < 0) {
            WN_PRINTF("ERROR YOU BLEW YOUR STACK\n");

            return;
        }

        WN_POINTER* retptrloc = &tlRetPtr;
        WN_POINTER* oldstackloc = &tlRetStackPtr;

        BeginFiber((WN_POINTER)retptrloc, (WN_POINTER)oldstackloc, s->startVal, (WN_POINTER)func, (WN_POINTER)stackloc, s->privateData->stacksize, stackleft);
    #endif
}

WN_VOID WNFiberManagerImpl::RunNextStack(WN_BOOL continuous) {
    tlContinuousRunning = continuous;

    WNFiberData* i = GetNextStack();

    if (i != WN_NULL) {
        #ifndef _WN_FIBER_USE_LIBRARIES
            tlCurStack = i;
        #endif

        ContinueOther(i);

        #ifdef _WN_FIBER_USE_LIBRARIES
            tlLastStack->privateData->cont = tlLastStackCont;
        #endif
    }
}

WN_BOOL WNFiberManagerImpl::UpdateQueues(WN_UINT32 currentTime) {
    if (WNAtomicCompareSwap(&mThreadUpdating, 1, 0) != 0) {
        return(WN_FALSE);
    }

    for (WN_UINT32 i = 0; i < mThreads.size();) {
        if (mThreads[i]->sleeptype == WN_SLP_INVALID) {
            WNFiberData* si = mThreads[i];

            mThreadsLock.Lock();
            mThreads[i] = mThreads[mThreads.size() - 1];
            mThreads.pop_back();
            mNumActive--;
            mThreadsLock.Unlock();

            DestroyFiberData(si);

            continue;
        } else if (mThreads[i]->sleeptype == WN_SLP_FINISHED) {
            WNFiberData* inf = mThreads[i];
            
            mThreadsLock.Lock();
            mThreads[i] = mThreads[mThreads.size() - 1];
            mThreads.pop_back();
            mNumActive--;
            mThreadsLock.Unlock();

            inf->sleeptype = WN_SLP_FINISHED;

            mDoneLock.Lock();
            mDoneThreads.push_back(inf);
            mDoneLock.Unlock();

            continue;
        } else if (mThreads[i]->runnable != 0) {
            if (mThreads[i]->sleeptype == WN_SLP_NOT_STARTED || mThreads[i]->sleeptype == WN_SLP_NO_SLEEP) {
                WNFiberData* inf = mThreads[i];

                inf->sleeptype = WN_SLP_READY;

                mRunnableLock.Lock();
                mRunnableThreads.push_back(inf);
                mRunnableLock.Unlock();

                continue;
            } else if (mThreads[i]->sleeptype == WN_SLP_SLEEP_TIME && (mThreads[i]->sleepval < currentTime)) {
                WNFiberData* inf = mThreads[i];

                inf->sleeptype = WN_SLP_READY;

                mRunnableLock.Lock();
                mRunnableThreads.push_back(inf);
                mRunnableLock.Unlock();

                continue;
            }
        }

        ++i;
    }

    mThreadUpdating = 0;

    return(WN_TRUE);
}

WN_BOOL WNFiberManagerImpl::QueuedThreads() {
    return(mNumActive > 0);
}

WN_UINT32 WNFiberManagerImpl::GetTime() {
    if (mTimeFunction) {
        return(mTimeFunction());
    }

    return(WNGetTickCount());
}
